package com.framework.bdf4j.msg.config;

import java.util.Map;

import javax.servlet.http.HttpSession;

import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.http.server.ServletServerHttpRequest;
import org.springframework.messaging.Message;
import org.springframework.messaging.MessageChannel;
import org.springframework.messaging.simp.config.ChannelRegistration;
import org.springframework.messaging.simp.config.MessageBrokerRegistry;
import org.springframework.messaging.simp.stomp.StompCommand;
import org.springframework.messaging.simp.stomp.StompHeaderAccessor;
import org.springframework.messaging.support.ChannelInterceptor;
import org.springframework.messaging.support.MessageHeaderAccessor;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.config.annotation.EnableWebSocketMessageBroker;
import org.springframework.web.socket.config.annotation.StompEndpointRegistry;
import org.springframework.web.socket.config.annotation.WebSocketMessageBrokerConfigurer;
import org.springframework.web.socket.server.HandshakeInterceptor;

/**
* Websocket 配置类
*
* @author shan.deng
* @date 2020-07-15 16:06:44
*/
@Configuration
@EnableWebSocketMessageBroker
public class WebSocketConfiguration  implements WebSocketMessageBrokerConfigurer {

    /**
    * 注册STOMP协议的节点（endpoint）并映射指定的URL
    * @param registry
    */
   public void registerStompEndpoints(StompEndpointRegistry registry){

       registry.addEndpoint("/net").addInterceptors(new HandshakeInterceptor() {
           @Override
           public boolean beforeHandshake(ServerHttpRequest serverHttpRequest, ServerHttpResponse serverHttpResponse, WebSocketHandler webSocketHandler, Map<String, Object> map) throws Exception {
               if (serverHttpRequest instanceof ServletServerHttpRequest) {
                   ServletServerHttpRequest servletRequest = (ServletServerHttpRequest) serverHttpRequest;
                   HttpSession session = servletRequest.getServletRequest().getSession();
                   map.put("sessionId", session.getId());
                   map.put("CURRENT_USERNAME",session.getAttribute("CURRENT_USERNAME").toString());
               }
               return true;
           }

           @Override
           public void afterHandshake(ServerHttpRequest serverHttpRequest, ServerHttpResponse serverHttpResponse, WebSocketHandler webSocketHandler, Exception e) {

           }
       }).setAllowedOrigins("*").withSockJS();
   }


   /**
    * 配置消息代理类
    * @param registry
    */
   public void configureMessageBroker(MessageBrokerRegistry registry){
       //广播式应配置一个消息结果消息代理
       registry.enableSimpleBroker("/net/result","/result");
       registry.setUserDestinationPrefix("/net");
   }

    @Bean
    public UserInterceptor createUserInterceptor() {
        return new UserInterceptor();
    }

   @Override
   public void configureClientInboundChannel(ChannelRegistration registration) {

       //注册一个拦截器用来做权限校验
       registration.interceptors(new ChannelInterceptor() {
           @Override
           public Message<?> preSend(Message<?> message, MessageChannel channel) {

               StompHeaderAccessor accessor = MessageHeaderAccessor.getAccessor(message, StompHeaderAccessor.class);
               if(StompCommand.CONNECT.equals(accessor.getCommand())) {
                   Map<String, Object> sessionAttributes = accessor.getSessionAttributes();
                   accessor.setUser(new UserPrincipal(sessionAttributes.get("CURRENT_USERNAME").toString()));
               }
               return message;
           }
       });
   }
}
